{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class MultiTaskModel(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, task1_output, task2_output):\n",
    "        super(MultiTaskModel, self).__init__()\n",
    "        self.shared_layer = nn.Linear(input_size, hidden_size)\n",
    "        self.task1_head = nn.Linear(hidden_size, task1_output)  # 任务1输出头\n",
    "        self.task2_head = nn.Linear(hidden_size, task2_output)  # 任务2输出头\n",
    "    \n",
    "    def forward(self, x):\n",
    "        shared_rep = torch.relu(self.shared_layer(x))\n",
    "        task1_output = self.task1_head(shared_rep)  # 任务1的结果\n",
    "        task2_output = self.task2_head(shared_rep)  # 任务2的结果\n",
    "        return task1_output, task2_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class MultiTaskModel(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, task1_output, task2_output):\n",
    "        super(MultiTaskModel, self).__init__()\n",
    "        self.shared_layer = nn.Linear(input_size, hidden_size)\n",
    "        self.task1_head = nn.Linear(hidden_size, task1_output)  # 任务1输出头\n",
    "        self.task2_head = nn.Linear(hidden_size, task2_output)  # 任务2输出头\n",
    "    \n",
    "    def forward(self, x):\n",
    "        shared_rep = torch.relu(self.shared_layer(x))\n",
    "        task1_output = self.task1_head(shared_rep)  # 任务1的结果\n",
    "        task2_output = self.task2_head(shared_rep)  # 任务2的结果\n",
    "        return task1_output, task2_output"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
